from langchain_community.document_loaders.csv_loader import CSVLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import OllamaEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langdev_helper.llm.qwen import llm

loader = CSVLoader(file_path='../../data/ordersample.csv')
data = loader.load()

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000, chunk_overlap=200, add_start_index=True
)
all_splits = text_splitter.split_documents(data)

embedding_model = OllamaEmbeddings(
    base_url='http://localhost:11434',
    model='quentinz/bge-large-zh-v1.5:q4_0'
)
vector = FAISS.from_documents(all_splits, embedding_model)

retriever = vector.as_retriever(search_type="similarity", search_kwargs={"k": 3})
# docs = retriever.invoke("收货人姓名是张三丰的，有几个订单？金额分别是多少，总共是多少？")
# docs = retriever.invoke("收货地址是朝阳区的有哪些订单？")
# for doc in docs:
    # print(doc)


prompt = ChatPromptTemplate.from_template("""仅根据所提供的上下文回答以下问题:

<context>
{context}
</context>

问题: {question}""")

from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
retriever_chain = (
    {"context": retriever , "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

# print(retriever_chain.invoke("订单ID是123456的收货人是谁，电话是多少?"))
# print(retriever_chain.invoke("收货人张三丰有几个订单？金额分别是多少，总共是多少？"))
print(retriever_chain.invoke("收货地址是朝阳区的有哪些订单？"))
